import os
import shutil
from tqdm import tqdm
import argparse
from typing import List
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime

def classfiy_data():
    '''
    classify data to 【prohibitory、warning、mandatory】
    '''
    labels_folder = r'D:\CCTSDB2021\test\labels'
    images_folder = r'D:\CCTSDB2021\test\images'
    warning_labels_folder = r'D:\CCTSDB2021'
    warning_images_folder = r'D:\CCTSDB2021'
    class_categories=['mandatory','prohibitory','warning']

    # 创建warning文件夹（如果不存在）
    for item in class_categories:
        os.makedirs(f'{warning_labels_folder}\\{item}\\eval\\labels', exist_ok=True)
        os.makedirs(f'{warning_labels_folder}\\{item}\\eval\\images', exist_ok=True)

    # 遍历labels文件夹
    for filename in os.listdir(labels_folder):
        label_file_path = os.path.join(labels_folder, filename)
        image_file_path = os.path.join(images_folder, filename.replace('.txt', '.jpg'))
        
        with open(label_file_path, 'r') as label_file:
            file_line = label_file.readlines()
            # if len(file_line) == 1:
            #     for every_row in file_line:
            #         if every_row.split(' ')[0] == '0':
            #             # 复制label文件到warning文件夹
            #             shutil.copy(label_file_path, os.path.join(f'{warning_labels_folder}\\mandatory\\eval\labels', filename))
            #             # 复制对应的图片文件到warning文件夹
            #             shutil.copy(image_file_path, os.path.join(f'{warning_labels_folder}\\mandatory\\eval\images', filename.replace('.txt', '.jpg')))
            #         elif every_row.split(' ')[0] == '1':
            #             shutil.copy(label_file_path, os.path.join(f'{warning_labels_folder}\\warning\\eval\labels', filename))
            #             # 复制对应的图片文件到warning文件夹
            #             shutil.copy(image_file_path, os.path.join(f'{warning_labels_folder}\\warning\\eval\images', filename.replace('.txt', '.jpg')))
            #         else:
            #             shutil.copy(label_file_path, os.path.join(f'{warning_labels_folder}\\prohibitory\\eval\labels', filename))
            #             # 复制对应的图片文件到warning文件夹
            #             shutil.copy(image_file_path, os.path.join(f'{warning_labels_folder}\\prohibitory\\eval\images', filename.replace('.txt', '.jpg')))
            # if len(file_line) > 1:

if __name__ =='__main__':
    classfiy_data()